package com.quanxiaoha.ai.robot.controller;

import com.quanxiaoha.ai.robot.model.AIResponse;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.http.MediaType;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Caius
 * @description
 * @since Created in 2025-09-07
 */
@RestController
@RequestMapping("/v6/ai")
public class AliyunBaiLianController {

    @Resource
    private OpenAiChatModel chatModel;

    // 储存聊天对话
    private Map<String, List<Message>> chatMemoryStore = new ConcurrentHashMap<>();

    /**
     * 普通对话
     * @param message
     * @return
     */
    @GetMapping("/generate")
    public String generate(@RequestParam(value = "message", defaultValue = "你是谁?") String message,
                           @RequestParam(value = "chatId") String chatId) {
        // 获取之前存在的记录
        List<Message> messages = chatMemoryStore.get(chatId);
        // 没有就初始化一份
        if (CollectionUtils.isEmpty(messages)) {
            messages = new ArrayList<>();
            chatMemoryStore.put(chatId, messages);
        }
        // 添加用户角色
        messages.add(new UserMessage(message));

        Prompt prompt = new Prompt(messages);

        String responseText = chatModel.call(prompt).getResult().getOutput().getText();

        // 添加 “助手角色消息” 到聊天记录中
        messages.add(new AssistantMessage(responseText));
        return responseText;
    }

    /**
     * 流式输出
     * @param message
     * @return
     */
    @GetMapping(value = "/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<AIResponse> generateStream(@RequestParam(value = "message", defaultValue = "你是谁?") String message) {
        // 系统角色
        SystemMessage systemMessage = new SystemMessage("请你扮演java面试官");
        // 用户角色
        UserMessage userMessage = new UserMessage(message);
        // 构建提示词
        Prompt prompt = new Prompt(Arrays.asList(systemMessage, userMessage));

        return chatModel.stream(prompt)
                .mapNotNull(chatResponse -> {
                    Generation generation = chatResponse.getResult();
                    String text = generation.getOutput().getText();
                    return AIResponse.builder().v(text).build();
                });
    }
}
